#include <boost/foreach.hpp>

#include "stl-helper.h"
#include "gene_processing.h"

MotifStats::MotifStats(boost::program_options::variables_map& vm, const Sequences& seqs)
  : vm_(vm), seqs_(seqs), acgt("ACGT")
{
  motif_ = vm["pattern"].as<std::string>();
  outputFile_ = vm["output"].as<std::string>();
}

void MotifStats::run()
{
  reset_stats();
  BOOST_FOREACH(const std::string& seq_i, seqs_) {
    vi_t pos;
    Sequences::check_occurence(motif_, seq_i, pos);
    BOOST_FOREACH(int p, pos) {
      count_stats(seq_i.substr(p+sz(motif_)));
    }
  }
  background_freq_ /= arma::sum(background_freq_);
  length_freq_ /= arma::sum(length_freq_);
  BOOST_FOREACH(vd_t& v, base_pos_freq_)
    v /= arma::sum(v);
  BOOST_FOREACH(vd_t& v, pos_base_freq_)
    v /= arma::sum(v);
}

void MotifStats::output()
{
  std::ofstream out;
  FileAccess::open2write(out, outputFile_);

  out << "base\t frequency" << ENDL;
  int c = 0;
  BOOST_FOREACH(double freq, background_freq_) {
    out << acgt[c++] << "\t " << freq << ENDL;
  }
  
  out << "length\t spacer_freq\t pos\t A\t C\t G\t T\t pos\t A\t C\t G\t T" << ENDL;
  int len = 0;
  BOOST_FOREACH(double freq, length_freq_) {
    out << len << "\t " << freq 
	<< "\t " << -len-1 << "\t" 
	<< base_pos_freq_[0][len] << "\t" << base_pos_freq_[1][len]  
	<< "\t" << base_pos_freq_[2][len] << "\t" << base_pos_freq_[3][len] 
	<< "\t " << -len-1 << "\t "
	<< "\t " << pos_base_freq_[len][0] << "\t " << pos_base_freq_[len][1] 
	<< "\t " << pos_base_freq_[len][2] << "\t " << pos_base_freq_[len][3] << ENDL;
    len++;
  }

  out.close();
}

void MotifStats::count_stats(const std::string& spacer)
{
  count_ += 1;
  length_freq_[sz(spacer)] += 1;
  int pos = -sz(spacer);
  BOOST_FOREACH(char c, spacer) {
    int c_pos = acgt.find(c);
    background_freq_[c_pos] += 1;
    base_pos_freq_[c_pos][-pos-1] += 1;
    pos_base_freq_[-pos-1][c_pos] += 1;
    pos++;
  }
}

void MotifStats::reset_stats()
{
  int max_length = seqs_.max_length();
  background_freq_ = arma::zeros<vd_t>(4);
  length_freq_ = arma::zeros(max_length+1);
  base_pos_freq_ = vvd_t(4, arma::zeros<vd_t>(max_length+1));
  pos_base_freq_ = vvd_t(max_length+1, arma::zeros<vd_t>(4));
  count_ = 0;
}

